import time
import uuid

from starlette.middleware.base import BaseHTTPMiddleware
from starlette.middleware.base import RequestResponseEndpoint
from starlette.requests import Request
from starlette.responses import Response

from rosa.scarborough_fair.generator_setting.settings import logger
from rosa.utils import REQUEST_CONTEXT


class RecordMiddleware(BaseHTTPMiddleware):
    async def dispatch(
            self, request: Request, call_next: RequestResponseEndpoint
    ) -> Response:
        # 生成唯一的请求ID
        request_id = request.headers.get("X-Request-ID", None)
        if not request_id:
            request_id = str(uuid.uuid4())
        request_id_var = REQUEST_CONTEXT["request_id"]
        request_id_var.set(request_id)
        start_time = time.time()
        message = f"Request ID: {request_id}, URL: {request.url}, Method: {request.method}, Headers: {request.headers}"
        logger.debug(message)
        response = await call_next(request)
        process_time = time.time() - start_time
        # 记录请求的响应信息
        message = f"Response ID: {request_id}, Status Code: {response.status_code}, Headers: {response.headers}, Process Time: {process_time}"
        logger.debug(message)
        return response


# 记录请求中间件
middleware = [{"middleware_class": RecordMiddleware}]
